package com.guoxin.tree.readblacktree;

/**
 * 树的骨架类
 * <p>
 * 红-黑规则：
 * 1.每个节点不是红色就是黑色的；
 * 2.根节点总是黑色的；
 * 3.如果节点是红色的，则它的子节点必须是黑色的（反之不一定）,(也就是从每个叶子到根的所有路径上不能有两个连续的红色节点)；
 * 4.从根节点到叶节点或空子节点的每条路径，必须包含相同数目的黑色节点（即相同的黑色高度）
 *
 * @author guoxin
 * @date 2019/7/28
 * @see <a href="https://www.cnblogs.com/ysocean/">参考博客</a>
 */
public class RBTree<T extends Comparable<T>> {

    /**
     * 红色节点标识
     */
    private static final boolean RED = true;

    /**
     * 黑色节点标识
     */
    private static final boolean BLACK = false;

    private Node<T> root;

    /**
     * 添加元素
     */
    public void add(Node<T> node) {
        int cmp;
        Node<T> root = this.root;
        Node<T> parent = null;

        //定位节点添加到哪个父节点下
        while (root != null) {
            parent = root;
            cmp = node.key.compareTo(root.key);
            if (cmp < 0) {
                root = root.left;
            } else {
                root = root.right;
            }
        }

        node.parent = parent;
        //表示当前没一个节点，那么就当新增的节点为根节点
        if (parent == null) {
            this.root = node;
        } else {
            //找出在当前父节点下新增节点的位置
            cmp = node.key.compareTo(parent.key);
            if (cmp < 0) {
                parent.left = node;
            } else {
                parent.right = node;
            }
        }
        //设置插入节点的颜色为红色
        node.color = RED;
        //修正为红黑树
        insertFixUp(node);
    }

    /**
     * 获取当前节点的父节点
     */
    private Node<T> getParent(Node<T> node) {
        return node == null ? null : node.parent;
    }

    /**
     * 设置当前节点为指定颜色
     */
    private void setColor(Node<T> node, boolean color) {
        if (node != null) {
            node.color = color;
        }
    }

    /**
     * 设置当前节点颜色为红色
     */
    private void setRedColor(Node<T> node) {
        if (node != null) {
            setColor(node, RED);
        }
    }

    /**
     * 设置当前节点颜色为黑色
     */
    private void setBlackColor(Node<T> node) {
        if (node != null) {
            setColor(node, BLACK);
        }
    }

    /**
     * 修正红黑树
     * <p>
     * 如果插入节点的父节点是黑色的，那不会违背红-黑树的规则，什么也不需要做
     * 但是遇到如下三种情况，我们就要开始变色和旋转了：
     * 1.插入节点的父节点和其叔叔节点（祖父节点的另一个子节点）均为红色。
     * 2.插入节点的父节点是红色的，叔叔节点是黑色的，且插入节点是其父节点的右子节点。
     * 3.插入节点的父节点是红色的，叔叔节点是黑色的，且插入节点是其父节点的左子节点。
     */
    private void insertFixUp(Node<T> node) {
        Node<T> parent, grandpa;

        //需要修正的条件：父节点存在，且父节点的颜色是红色
        while (((parent = node.parent) != null) && parent.color == RED) {
            grandpa = getParent(parent);

            //父节点是祖父节点的左孩子
            if (parent == grandpa.left) {
                Node<T> uncle = grandpa.right;//叔叔节点

                //case1:叔叔节点也是红色
                if (uncle.color == RED) {
                    setBlackColor(uncle);
                    setBlackColor(parent);
                    setRedColor(grandpa);
                    node = grandpa;
                    continue;
                }

                //case2:叔叔节点是黑色，且当前节点是右子节点
                if (node == parent.right) {
                    leftRotate(parent);//从父节点出左旋
                    Node<T> tmp = parent;//然后将父节点和自己调换一下，为下面右旋做准备
                    parent = node;
                    node = tmp;
                }

                //case3:叔叔节点是黑色，且当前节点是左子节点
                setBlackColor(parent);
                setRedColor(grandpa);
                rightRotate(grandpa);
            } else {//父节点是祖父节点的右孩子
                Node<T> uncle = grandpa.left;

                //case1:叔叔节点也是红色的
                if (uncle.color == RED) {
                    setBlackColor(parent);
                    setBlackColor(uncle);
                    setRedColor(grandpa);
                    node = grandpa;
                    continue;
                }

                //case2:叔叔节点是黑色的，且当前节点是左子节点
                if (node == parent.left) {
                    rightRotate(parent);
                    Node<T> tmp = parent;//然后将父节点和自己调换一下，为下面右旋做准备
                    parent = node;
                    node = tmp;
                }

                //case3:叔叔节点是黑色的，且当前节点是右子节点
                setBlackColor(parent);
                setRedColor(grandpa);
                leftRotate(grandpa);
            }
        }
        setBlackColor(node);
    }

    /**
     * 左旋
     * <p>
     * 左旋做了三件事：
     * 1. 将y的左子节点赋给x的右子节点,并将x赋给y左子节点的父节点(y左子节点非空时)
     * 2. 将x的父节点p(非空时)赋给y的父节点，同时更新p的子节点为y(左或右)
     * 3. 将y的左子节点设为x，将x的父节点设为y
     */
    @SuppressWarnings("all")
    private void leftRotate(Node<T> x) {
        //     p                       p
        //     /                       /
        //    x                       y
        //   / \                     / \
        //  lx  y      ----->      x  ry
        //     / \                 / \
        //    ly ry               lx ly

        //1. 将y的左子节点赋给x的右子节点，并将x赋给y左子节点的父节点(y左子节点非空时)
        Node<T> y = x.right;
        x.right = y.left;
        if (y.left != null) {
            y.left.parent = x;
        }

        //2. 将x的父节点p(非空时)赋给y的父节点，同时更新p的子节点为y(左或右)
        y.parent = x.parent;
        if (x.parent == null) {
            this.root = y; //如果x的父节点为空(即x为根节点)，则将y设为根节点
        } else {
            //如果x是左子节点
            if (x == x.parent.left) {
                x.parent.left = y;//则也将y设为左子节点
            } else {
                x.parent.right = y;//否则将y设为右子节点
            }
        }

        //3. 将y的左子节点设为x，将x的父节点设为y
        y.left = x;
        x.parent = y;
    }


    /**
     * 右旋
     * <p>
     * 右旋做了三件事：
     * 1. 将x的右子节点赋给y的左子节点,并将y赋给x右子节点的父节点(x右子节点非空时)
     * 2. 将y的父节点p(非空时)赋给x的父节点，同时更新p的子节点为x(左或右)
     * 3. 将x的右子节点设为y，将y的父节点设为x
     */
    @SuppressWarnings("all")
    private void rightRotate(Node<T> y) {
        //        p                   p
        //       /                   /
        //      y                   x
        //     / \                 / \
        //    x  ry   ----->      lx  y
        //   / \                     / \
        // lx  rx                   rx ry

        //1. 将y的左子节点赋给x的右子节点，并将x赋给y左子节点的父节点(y左子节点非空时)
        Node<T> x = y.left;
        y.left = x.right;
        if (x.right != null) {
            x.right.parent = y;
        }

        //2. 将x的父节点p(非空时)赋给y的父节点，同时更新p的子节点为y(左或右)
        x.parent = y.parent;
        if (y.parent == null) {
            this.root = x;//如果y的父节点为空(即y为根节点)，则旋转后将x设为根节点
        } else {
            //如果y是左子节点
            if (y == y.parent.left) {
                y.parent.left = x;//则将x也设置为左子节点
            } else {
                y.parent.right = x; //否则将x设置为右子节点
            }
        }

        //3. 将x的左子节点设为y，将y的父节点设为y
        x.right = y;
        y.parent = x;
    }


    private static class Node<T extends Comparable<T>> {

        /**
         * 用于比较的key值
         */
        T key;
        /**
         * 指向左孩子节点的引用
         */
        Node<T> left;
        /**
         * 指向右孩子节点的引用
         */
        Node<T> right;
        /**
         * 指向父节点的引用
         */
        Node<T> parent;

        /**
         * 颜色 true表示黑，false表示红
         */
        boolean color;

        public Node(T key, Node<T> left, Node<T> right, Node<T> parent, boolean color) {
            this.key = key;
            this.left = left;
            this.right = right;
            this.parent = parent;
            this.color = color;
        }
    }
}
